package de.dailab.irml.hb;

import java.io.File;
import java.io.FileOutputStream;
import java.io.OutputStreamWriter;
import java.util.Map;
import java.util.Random;

import org.jblas.DoubleMatrix;

import de.dailab.irml.ml.classification.ANNClassifier;
import de.dailab.irml.ml.core.PSGD;
import de.dailab.irml.ml.evaluation.AMSMetric;
import de.dailab.irml.ml.evaluation.IntegratedEvaluation;
import de.dailab.irml.ml.evaluation.RMSE;
import de.dailab.irml.ml.feature.NormalScaler;

import weka.core.Attribute;
import weka.core.Instance;
import weka.core.Instances;
import weka.core.converters.CSVLoader;

public class HiggsBosonPTest {
	private double trainingRatio;
	private String inputPath;
	private DoubleMatrix x;
	private DoubleMatrix y;
	private DoubleMatrix c;
	private DoubleMatrix x_p;
	private DoubleMatrix y_p;

	public DoubleMatrix getW_p() {
		return w_p;
	}

	public void setW_p(DoubleMatrix w_p) {
		this.w_p = w_p;
	}

	private DoubleMatrix w_p;
	private DoubleMatrix w;

	public DoubleMatrix getW() {
		return w;
	}

	public void setW(DoubleMatrix w) {
		this.w = w;
	}

	public HiggsBosonPTest(String inputPath, double trainingRatio) throws Exception {
		this.trainingRatio = trainingRatio;
		this.inputPath = inputPath;
	}

	public void loadTraining() throws Exception {
		System.out.println("loading training");
		File f = new File(inputPath);
		CSVLoader loader = new CSVLoader();
		loader.setSource(f);

		System.out.println("converting training");

		Instances input = loader.getDataSet();

		System.out.println("spliting data with training " + trainingRatio + "%");
		input.randomize(new Random(1));
		int trainSize = (int) Math.round(input.numInstances() * (trainingRatio / 100));
		int testSize = input.numInstances() - trainSize;
		Instances training = new Instances(input, 0, trainSize);
		Instances test = new Instances(input, trainSize, testSize);
		input = training;
		Instances output = test;

		double maxP = 0;
		double maxN = 0;
		for (int i = 0; i < input.size(); i++) {
			Instance dp = input.get(i);
			double w = dp.value(31);
			double c = dp.value(32);
			for (int j = 1; j <= 30; j++) {
				if (dp.value(j) < -998)
					dp.setMissing(j);
			}
			if (c == 1) {
				if (w > maxN)
					maxN = w;
			} else {
				if (w > maxP)
					maxP = w;
			}
		}

		// maxP = maxP / 2;

		for (int i = 0; i < input.size(); i++) {
			Instance dp = input.get(i);
			double c = dp.value(32);
			double w = dp.value(31);
			dp.setValue(31, Math.log(w));
			if (c == 1) {
				w = w / maxN;
			} else {
				w = w / maxP;
			}
			dp.setWeight(w);
		}

		for (int i = 0; i < output.size(); i++) {
			Instance dp = output.get(i);
			for (int j = 1; j <= 30; j++) {
				if (dp.value(j) < -998)
					dp.setMissing(j);
			}
		}
		System.out.println("training converted, size: " + input.size());
		System.out.println("test converted, size: " + output.size());
	}

	public DoubleMatrix getX() {
		return x;
	}

	public void setX(DoubleMatrix x) {
		this.x = x;
	}

	public DoubleMatrix getY() {
		return y;
	}

	public void setY(DoubleMatrix y) {
		this.y = y;
	}

	public DoubleMatrix getC() {
		return c;
	}

	public void setC(DoubleMatrix c) {
		this.c = c;
	}

	public DoubleMatrix getX_p() {
		return x_p;
	}

	public void setX_p(DoubleMatrix x_p) {
		this.x_p = x_p;
	}

	public DoubleMatrix getY_p() {
		return y_p;
	}

	public void setY_p(DoubleMatrix y_p) {
		this.y_p = y_p;
	}

	public void loadTrainingReplaceMissing() throws Exception {
		System.out.println("loading training");
		File f = new File(inputPath);
		CSVLoader loader = new CSVLoader();
		loader.setSource(f);

		System.out.println("converting training");
		Instances input = loader.getDataSet();

		for (int k = 0; k < 30; k++) {
			Attribute atr = new Attribute(input.attribute(k + 1).name() + "_missing");
			input.insertAttributeAt(atr, 33 + k);
		}

		double maxP = 0;
		double maxN = 0;
		for (int i = 0; i < input.size(); i++) {
			Instance dp = input.get(i);
			double w = dp.value(31);
			String c = dp.stringValue(32);
			for (int j = 1; j <= 30; j++) {
				if (dp.value(j) < -998) {
					dp.setValue(j, 0);
					dp.setValue(32 + j, 1);
				} else {
					dp.setValue(32 + j, 0);
				}
			}
			if (c.equals("b")) {
				if (w > maxN)
					maxN = w;
			} else {
				if (w > maxP)
					maxP = w;
			}
		}

		for (int i = 0; i < input.size(); i++) {
			Instance dp = input.get(i);
			String c = dp.stringValue(32);
			double w = dp.value(31);
			if (c.equals("b")) {
				w = w / maxN;
			} else {
				w = w / maxP;
			}
			dp.setWeight(w);
		}

		System.out.println("spliting data with training " + trainingRatio + "%");
		input.randomize(new Random(1));
		int trainSize = (int) Math.round(input.numInstances() * (trainingRatio / 100));
		int testSize = input.numInstances() - trainSize;
		Instances training = new Instances(input, 0, trainSize);
		Instances test = new Instances(input, trainSize, testSize);
		input = training;
		Instances output = test;

		System.out.println("training converted, size: " + input.size());
		System.out.println("test converted, size: " + output.size());
		x = new DoubleMatrix(input.numInstances(), input.numAttributes() - 3);
		y = new DoubleMatrix(input.numInstances(), 1);
		c = new DoubleMatrix(input.numInstances(), 1);
		w = new DoubleMatrix(input.numInstances(), 1);

		x_p = new DoubleMatrix(output.numInstances(), output.numAttributes() - 3);
		y_p = new DoubleMatrix(output.numInstances(), 1);
		w_p = new DoubleMatrix(output.numInstances(), 1);
		for (int i = 0; i < input.numInstances(); i++) {
			Instance dp = input.get(i);
			// c.put(i, 0, 1.0);
			c.put(i, 0, dp.weight());
			for (int j = 1; j < input.numAttributes(); j++) {
				if (j == 31) {
					w.put(i, 0, dp.value(j));
				} else if (j == 32) {
					if (dp.stringValue(32).equals("b"))
						y.put(i, 0, 0.0);
					else
						y.put(i, 0, 1.0);
				} else {
					if (j >= 31)
						x.put(i, j - 3, dp.value(j));
					else
						x.put(i, j - 1, dp.value(j));
				}
			}
		}

		for (int i = 0; i < output.numInstances(); i++) {
			Instance dp = output.get(i);
			for (int j = 1; j < output.numAttributes(); j++) {
				if (j == 31) {
					w_p.put(i, 0, dp.value(j));
				} else if (j == 32) {
					if (dp.stringValue(32).equals("b"))
						y_p.put(i, 0, 0.0);
					else
						y_p.put(i, 0, 1.0);
				} else {
					if (j >= 31)
						x_p.put(i, j - 3, dp.value(j));
					else
						x_p.put(i, j - 1, dp.value(j));
				}
			}
		}

	}

	public static void main(String[] args) {
		try {
			Double r = Double.valueOf(args[1]);
			Integer numThreads = Integer.valueOf(args[2]);
			Double stepSize = Double.valueOf(args[3]);
			Double l2 = Double.valueOf(args[4]);
			Integer steps = Integer.valueOf(args[5]);
			Integer ANNM = Integer.valueOf(args[6]);
			String fileName = args[7];
			HiggsBosonPTest hb;
			hb = new HiggsBosonPTest(args[0], r.doubleValue());
			hb.loadTrainingReplaceMissing();

			// normalization
			System.out.println("data normalization");
			NormalScaler norm = new NormalScaler(hb.x.columns);
			norm.addData("train", hb.getX());
			norm.addData("test", hb.getX_p());
			Map<String, DoubleMatrix> datasets = norm.normalize();
			DoubleMatrix x = datasets.get("train");
			DoubleMatrix x_p = datasets.get("test");
			System.out.println("data normalized");
			ANNClassifier clf = new ANNClassifier(hb.getY(), hb.getC(), x, ANNM, numThreads);
			clf.setL2(l2);
			PSGD psgd = new PSGD(clf, stepSize, steps, numThreads);
			psgd.setDebug(true);
			double finalLoss = psgd.optimise();
			AMSMetric ams = new AMSMetric();
			RMSE wnl = new RMSE();
			IntegratedEvaluation ie = new IntegratedEvaluation();
			ie.addMeasurement(ams);
			ie.addMeasurement(wnl);
			ie.setPredictor(clf);
			ie.test(hb.getY_p(), x_p, hb.getW_p());
			String sOut = "Final Training Loss: " + finalLoss + "\n";
			sOut = sOut + ie.evaluate();
			System.out.println(sOut);
			FileOutputStream fo = new FileOutputStream(fileName);
			OutputStreamWriter fw = new OutputStreamWriter(fo);
			fw.write(sOut);
			fw.close();
			fo.close();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}
}
